Allow flexible positions in Gemma #2389
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @abheesht17, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request enhances the Gemma model's capability to handle flexible positional embeddings. This is a crucial step towards supporting advanced training techniques such as DPO, GRPO, and PPO, which often necessitate the use of left-padded prompts. The changes involve introducing a new utility function to compute positions from padding masks and integrating this new mechanism into Gemma's attention and decoder blocks, ensuring the model can correctly process inputs with varied padding structures.
Highlights
- New utility for position computation: A
compute_positions_from_maskfunction was added totransformer_layer_utils.pyto derive token positions from a padding mask. - Gemma model adaptation for flexible positions: The
GemmaDecoderBlockandCachedGemmaAttentionlayers were updated to accept and utilize these flexible positions, specifically by passing them to the RoPE layer. - Test coverage for flexible positions: New tests were introduced in
gemma_backbone_test.pyand a helper methodrun_positions_testwas added totest_case.pyto validate the correct behavior of flexible positional embeddings.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request introduces support for flexible token positions in the Gemma model, which is a key step for handling left-padded inputs required by algorithms like DPO. The changes include a new utility function to compute positions from a mask, and modifications to the Gemma attention and decoder blocks to utilize these positions. The addition of run_positions_test is a great way to ensure this new functionality is correct.
My review has identified one critical issue that will cause a TypeError at runtime, and a couple of medium-severity suggestions to improve code clarity and efficiency. Please see the detailed comments below.
| corresponding to positions of tokens in the sequence. | ||
| """ | ||
| positions = ops.cumsum(mask, axis=-1) | ||
| positions = ops.subtract(positions, ops.greater_equal(positions, 1)) |
| @@ -720,6 +720,47 @@ def compare(actual, expected): | |||
| output = ops.argmax(output, axis=-1) | |||
| self.assertAllEqual(output, expected_labels) | |||
|
|
|||
| def run_positions_test( | |||
There was a problem hiding this comment.
do we think this is generic and will extend beyond gemma? if so ok to leave here. if not I might park this directly in the gemma tests.
There was a problem hiding this comment.
Yeah, we should do this for all CausalLMs. I've done it only for Gemma for now, will extend to other models later.
| @@ -720,6 +720,47 @@ def compare(actual, expected): | |||
| output = ops.argmax(output, axis=-1) | |||
| self.assertAllEqual(output, expected_labels) | |||
|
|
|||
| def run_positions_test( | |||
There was a problem hiding this comment.
Yeah, we should do this for all CausalLMs. I've done it only for Gemma for now, will extend to other models later.
Follows #2369
We left-pad prompts for DPO (and GRPO/PPO), in which case, we need to compute flexible positions. This PR does it for Gemma, but we can extend it to other models soon.
Delta
GemmaDecoderBlock,CachedGemmaAttentionin order to pass these flexible positions.This PR does not, however, modify the
generate()code to account for left-padded prompts. This will be a more involved changes, with modifications required togenerate_preprocess(); it makes sense to do this later when we start work on GRPO/PPO.